#include "ScreenCapture.h"
#include "GDICapture.h"
#include "DXGICapture.h"
#include "DesktopMonitor.h"
#include "BMPSaver.h"
#include "SecureDesktopCapture.h"
#include "SystemServiceCapture.h"
#include "Utils.h"
#include <iostream>

ScreenCapture::ScreenCapture() :
    m_activeMethod(CaptureMethod::AUTO),
    m_isCapturing(false),
    m_shouldStop(false),
    m_captureThread(NULL),
    m_stopEvent(NULL) {

    m_gdiCapture = std::make_unique<GDICapture>();
    m_dxgiCapture = std::make_unique<DXGICapture>();
    m_desktopMonitor = std::make_unique<DesktopMonitor>();
    m_bmpSaver = std::make_unique<BMPSaver>();
    m_secureCapture = std::make_unique<SecureDesktopCapture>();
    m_systemServiceCapture = std::make_unique<SystemServiceCapture>();
}

ScreenCapture::~ScreenCapture() {
    StopCapture();
}

CaptureResult ScreenCapture::Initialize(const CaptureSettings& settings) {
    Utils::DebugLog("ScreenCapture: Initializing...");

    m_settings = settings;

    if (!m_bmpSaver->CreateDirectoryIfNotExists(m_settings.outputPath)) {
        m_lastError = "Failed to create output directory: " + m_bmpSaver->GetLastError();
        return CaptureResult::ERROR_INIT_FAILED;
    }

    if (m_settings.enableDesktopMonitoring) {
        if (!m_desktopMonitor->Initialize()) {
            m_lastError = "Failed to initialize desktop monitor: " + m_desktopMonitor->GetLastError();
            return CaptureResult::ERROR_INIT_FAILED;
        }
    }

    CaptureResult result = InitializeCaptureMethods();
    if (result != CaptureResult::SUCCESS) {
        return result;
    }

    m_stopEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
    if (!m_stopEvent) {
        m_lastError = "Failed to create stop event: " + Utils::GetLastWin32Error();
        return CaptureResult::ERROR_INIT_FAILED;
    }

    Utils::DebugLog("ScreenCapture: Initialized successfully");
    return CaptureResult::SUCCESS;
}

CaptureResult ScreenCapture::InitializeCaptureMethods() {
    if (m_settings.method == CaptureMethod::AUTO || m_settings.method == CaptureMethod::DXGI) {
        if (m_dxgiCapture->IsSupported()) {
            if (m_dxgiCapture->Initialize()) {
                if (m_settings.method == CaptureMethod::AUTO) {
                    m_activeMethod = CaptureMethod::DXGI;
                    Utils::DebugLog("ScreenCapture: Auto-selected DXGI capture method");
                } else {
                    m_activeMethod = CaptureMethod::DXGI;
                }
            } else {
                Utils::DebugLog("ScreenCapture: DXGI initialization failed: " + m_dxgiCapture->GetLastError());
                if (m_settings.method == CaptureMethod::DXGI) {
                    m_lastError = "DXGI capture requested but failed to initialize: " + m_dxgiCapture->GetLastError();
                    return CaptureResult::ERROR_INIT_FAILED;
                }
            }
        }
    }

    if (m_activeMethod == CaptureMethod::AUTO || m_settings.method == CaptureMethod::GDI) {
        if (!m_gdiCapture->Initialize()) {
            m_lastError = "GDI capture failed to initialize: " + m_gdiCapture->GetLastError();
            return CaptureResult::ERROR_INIT_FAILED;
        }

        if (m_activeMethod == CaptureMethod::AUTO) {
            m_activeMethod = CaptureMethod::GDI;
            Utils::DebugLog("ScreenCapture: Auto-selected GDI capture method");
        } else if (m_settings.method == CaptureMethod::GDI) {
            m_activeMethod = CaptureMethod::GDI;
        }
    }

    return CaptureResult::SUCCESS;
}

CaptureResult ScreenCapture::StartCapture() {
    if (m_isCapturing) {
        m_lastError = "Capture already in progress";
        return CaptureResult::ERROR_INIT_FAILED;
    }

    Utils::DebugLog("ScreenCapture: Starting capture...");

    m_shouldStop = false;
    ResetEvent(m_stopEvent);

    m_captureThread = CreateThread(NULL, 0, CaptureThreadProc, this, 0, NULL);
    if (!m_captureThread) {
        m_lastError = "Failed to create capture thread: " + Utils::GetLastWin32Error();
        return CaptureResult::ERROR_INIT_FAILED;
    }

    m_isCapturing = true;
    Utils::DebugLog("ScreenCapture: Capture started successfully");
    return CaptureResult::SUCCESS;
}

CaptureResult ScreenCapture::StopCapture() {
    if (!m_isCapturing) {
        return CaptureResult::SUCCESS;
    }

    Utils::DebugLog("ScreenCapture: Stopping capture...");

    m_shouldStop = true;
    if (m_stopEvent) {
        SetEvent(m_stopEvent);
    }

    if (m_captureThread) {
        if (WaitForSingleObject(m_captureThread, 5000) == WAIT_TIMEOUT) {
            Utils::DebugLog("ScreenCapture: Force terminating capture thread");
            TerminateThread(m_captureThread, 0);
        }
        CloseHandle(m_captureThread);
        m_captureThread = NULL;
    }

    if (m_stopEvent) {
        CloseHandle(m_stopEvent);
        m_stopEvent = NULL;
    }

    m_isCapturing = false;
    Utils::DebugLog("ScreenCapture: Capture stopped");
    return CaptureResult::SUCCESS;
}

CaptureResult ScreenCapture::CaptureFrame() {
    if (m_settings.enableDesktopMonitoring && m_desktopMonitor->HasDesktopChanged()) {
        Utils::DebugLog("ScreenCapture: Desktop change detected, handling...");
        CaptureResult result = HandleDesktopSwitch();
        if (result != CaptureResult::SUCCESS) {
            return result;
        }
    }

    HBITMAP hBitmap = NULL;
    int width = 0, height = 0;
    bool captureSuccess = false;

    // On secure desktop, prioritize GDI method
    DesktopType currentDesktopType = m_desktopMonitor->GetCurrentDesktopType();
    bool isSecureDesktop = (currentDesktopType == DesktopType::SECURE ||
                           currentDesktopType == DesktopType::WINLOGON ||
                           currentDesktopType == DesktopType::UAC);

    if (isSecureDesktop) {
        Utils::DebugLog("ScreenCapture: Secure desktop detected - Using system service capture method");

        // Method 1: Try SystemServiceCapture (highest privilege method)
        if (!m_systemServiceCapture->Initialize()) {
            Utils::DebugLog("ScreenCapture: Failed to initialize system service capture: " + m_systemServiceCapture->GetLastError());
        } else {
            captureSuccess = m_systemServiceCapture->CaptureActiveSession(hBitmap, width, height);
            if (captureSuccess) {
                Utils::DebugLog("ScreenCapture: System service capture succeeded - captured actual secure desktop!");
            } else {
                Utils::DebugLog("ScreenCapture: System service capture failed: " + m_systemServiceCapture->GetLastError());
            }
        }

        // Method 2: Fallback to SecureDesktopCapture
        if (!captureSuccess) {
            Utils::DebugLog("ScreenCapture: Trying secure desktop capture as fallback");
            if (m_secureCapture->Initialize()) {
                captureSuccess = m_secureCapture->CaptureSecureDesktop(hBitmap, width, height);
                if (captureSuccess) {
                    Utils::DebugLog("ScreenCapture: Secure desktop capture succeeded");
                } else {
                    Utils::DebugLog("ScreenCapture: Secure desktop capture failed: " + m_secureCapture->GetLastError());
                }
            }
        }

        // Method 3: Last resort - try standard GDI with forced reinit
        if (!captureSuccess) {
            Utils::DebugLog("ScreenCapture: Trying GDI as last resort for secure desktop");
            m_gdiCapture->Cleanup();
            if (m_gdiCapture->Initialize()) {
                captureSuccess = m_gdiCapture->CaptureScreen(hBitmap, width, height);
                if (captureSuccess) {
                    Utils::DebugLog("ScreenCapture: GDI capture succeeded (content may be limited)");
                }
            }
        }
    } else {
        // Normal desktop, use original logic
        if (m_activeMethod == CaptureMethod::DXGI && m_dxgiCapture->IsInitialized()) {
            captureSuccess = m_dxgiCapture->CaptureScreen(hBitmap, width, height);
            if (!captureSuccess) {
                Utils::DebugLog("ScreenCapture: DXGI capture failed: " + m_dxgiCapture->GetLastError());
                if (SwitchCaptureMethod(CaptureMethod::GDI) == CaptureResult::SUCCESS) {
                    captureSuccess = m_gdiCapture->CaptureScreen(hBitmap, width, height);
                    if (!captureSuccess) {
                        Utils::DebugLog("ScreenCapture: GDI fallback failed: " + m_gdiCapture->GetLastError());
                    }
                }
            }
        } else if (m_activeMethod == CaptureMethod::GDI && m_gdiCapture->IsInitialized()) {
            captureSuccess = m_gdiCapture->CaptureScreen(hBitmap, width, height);
            if (!captureSuccess) {
                Utils::DebugLog("ScreenCapture: GDI capture failed: " + m_gdiCapture->GetLastError());
            }
        }
    }

    if (!captureSuccess) {
        std::string errorDetail = "Desktop type: ";
        switch (currentDesktopType) {
            case DesktopType::NORMAL: errorDetail += "Normal"; break;
            case DesktopType::SECURE: errorDetail += "Secure"; break;
            case DesktopType::WINLOGON: errorDetail += "Winlogon"; break;
            case DesktopType::UAC: errorDetail += "UAC"; break;
            case DesktopType::SCREENSAVER: errorDetail += "Screensaver"; break;
            default: errorDetail += "Unknown"; break;
        }
        m_lastError = "All capture methods failed on " + errorDetail;
        return CaptureResult::ERROR_CAPTURE_FAILED;
    }

    std::string filename = GenerateFilename();
    bool saveSuccess = m_bmpSaver->SaveBitmap(hBitmap, filename);

    DeleteObject(hBitmap);

    if (!saveSuccess) {
        m_lastError = "Failed to save bitmap: " + m_bmpSaver->GetLastError();
        return CaptureResult::ERROR_SAVE_FAILED;
    }

    return CaptureResult::SUCCESS;
}

CaptureResult ScreenCapture::SwitchCaptureMethod(CaptureMethod method) {
    Utils::DebugLog("ScreenCapture: Switching capture method");

    if (method == CaptureMethod::DXGI) {
        if (m_dxgiCapture->IsSupported() && !m_dxgiCapture->IsInitialized()) {
            if (m_dxgiCapture->Initialize()) {
                m_activeMethod = CaptureMethod::DXGI;
                return CaptureResult::SUCCESS;
            }
        }
    } else if (method == CaptureMethod::GDI) {
        if (!m_gdiCapture->IsInitialized()) {
            if (m_gdiCapture->Initialize()) {
                m_activeMethod = CaptureMethod::GDI;
                return CaptureResult::SUCCESS;
            }
        } else {
            m_activeMethod = CaptureMethod::GDI;
            return CaptureResult::SUCCESS;
        }
    }

    return CaptureResult::ERROR_INIT_FAILED;
}

CaptureResult ScreenCapture::HandleDesktopSwitch() {
    DesktopType desktopType = m_desktopMonitor->GetCurrentDesktopType();
    std::string desktopName = m_desktopMonitor->GetDesktopName();

    Utils::DebugLog("ScreenCapture: Desktop switch to: " + desktopName);

    if (!m_settings.captureSecureDesktops &&
        (desktopType == DesktopType::SECURE ||
         desktopType == DesktopType::WINLOGON ||
         desktopType == DesktopType::UAC)) {
        m_lastError = "Secure desktop detected but capture disabled by settings";
        return CaptureResult::ERROR_DESKTOP_SWITCH;
    }

    // Handle secure desktop specially
    bool isSecureDesktop = (desktopType == DesktopType::SECURE ||
                           desktopType == DesktopType::WINLOGON ||
                           desktopType == DesktopType::UAC);

    if (isSecureDesktop) {
        Utils::DebugLog("ScreenCapture: Handling secure desktop switch");

        // On secure desktop, DXGI usually doesn't work, disable it
        if (m_dxgiCapture->IsInitialized()) {
            m_dxgiCapture->Cleanup();
        }

        // Ensure GDI capture is available
        if (m_gdiCapture->IsInitialized()) {
            m_gdiCapture->Cleanup();
        }

        if (!m_gdiCapture->Initialize()) {
            m_lastError = "Failed to initialize GDI capture for secure desktop";
            return CaptureResult::ERROR_DESKTOP_SWITCH;
        }

        // Force GDI method on secure desktop
        m_activeMethod = CaptureMethod::GDI;
    } else {
        // Normal desktop, reinitialize all capture methods
        if (m_gdiCapture->IsInitialized()) {
            m_gdiCapture->Cleanup();
            if (!m_gdiCapture->Initialize()) {
                m_lastError = "Failed to reinitialize GDI capture after desktop switch";
                return CaptureResult::ERROR_DESKTOP_SWITCH;
            }
        }

        if (m_dxgiCapture->IsInitialized()) {
            m_dxgiCapture->Cleanup();
        }

        // If auto or DXGI mode, try to reinitialize DXGI
        if ((m_settings.method == CaptureMethod::AUTO || m_settings.method == CaptureMethod::DXGI) &&
            m_dxgiCapture->IsSupported()) {
            if (m_dxgiCapture->Initialize()) {
                m_activeMethod = CaptureMethod::DXGI;
            } else {
                m_activeMethod = CaptureMethod::GDI;
            }
        } else {
            m_activeMethod = CaptureMethod::GDI;
        }
    }

    return CaptureResult::SUCCESS;
}

std::string ScreenCapture::GenerateFilename() const {
    return Utils::GetFullOutputPath(m_settings.outputPath);
}

DWORD WINAPI ScreenCapture::CaptureThreadProc(LPVOID lpParam) {
    ScreenCapture* pThis = static_cast<ScreenCapture*>(lpParam);
    pThis->CaptureLoop();
    return 0;
}

void ScreenCapture::CaptureLoop() {
    Utils::DebugLog("ScreenCapture: Capture loop started");

    while (!m_shouldStop) {
        CaptureResult result = CaptureFrame();

        if (result != CaptureResult::SUCCESS) {
            Utils::DebugLog("ScreenCapture: Capture frame failed: " + m_lastError);
        }

        if (WaitForSingleObject(m_stopEvent, m_settings.captureInterval) == WAIT_OBJECT_0) {
            break;
        }
    }

    Utils::DebugLog("ScreenCapture: Capture loop ended");
}